#!/usr/bin/env python3

import socket
import subprocess
import os
import pty
import select
import sys
import termios
import struct
import fcntl
import argparse
import signal

def set_winsize(master_fd, row, col):
    """设置 PTY 窗口大小"""
    import struct
    winsize = struct.pack("HHHH", row, col, 0, 0)
    fcntl.ioctl(master_fd, termios.TIOCSWINSZ, winsize)

def get_winsize(master_fd):
    """获取当前 PTY 窗口大小"""
    try:
        winsize = fcntl.ioctl(master_fd, termios.TIOCGWINSZ, b'\x00' * 8)
        row, col = struct.unpack("HHHH", winsize)[:2]
        return row, col
    except:
        return 24, 80

def setup_pty():
    """子进程启动前配置终端属性"""
    try:
        # 获取 slave 端（0 是 stdin）
        attrs = termios.tcgetattr(0)
        # 设置回显和规范模式
        attrs[3] &= termios.ECHO      # 回显on
        attrs[3] &= ~termios.ICANON    # 原始模式
        attrs[3] &= ~termios.ECHOE     # 禁用退格回显
        attrs[3] &= ~termios.ECHOK     # 禁用换行回显
        attrs[6][termios.VMIN] = 1     # 最小读取字符
        attrs[6][termios.VTIME] = 0    # 超时
        termios.tcsetattr(0, termios.TCSANOW, attrs)
    except:
        pass

def handle_client(client_sock, client_addr, initial_cols=80, initial_rows=24):
    print(f"[+] 已接受来自 {client_addr[0]}:{client_addr[1]} 的连接")

    # 创建 PTY
    master, slave = pty.openpty()

    # 设置初始窗口大小
    set_winsize(master, initial_rows, initial_cols)

    # 设置环境变量，防止 Vim 输出复杂控制序列
    env = os.environ.copy()
    env["TERM"] = "xterm"        # 关键：避免颜色和图形终端序列，TERM=dumb
    env["PATH"] = "/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin"  # 确保 PATH 正常
    env["PS1"] = r"\u@\h:\w\$ "  # 可选：设置简单提示符

    try:
        proc = subprocess.Popen(
            ["/bin/bash"],
            stdin=slave,
            stdout=slave,
            stderr=slave,
            env=env,
            preexec_fn=setup_pty  # 子进程启动前配置终端
        )
    except Exception as e:
        print(f"[-] 启动 bash 失败: {e}")
        client_sock.close()
        os.close(master)
        os.close(slave)
        return

    os.close(slave)  # 关闭 slave 端（由子进程使用）

    # 设置 master 为非阻塞
    flags = fcntl.fcntl(master, fcntl.F_GETFL)
    fcntl.fcntl(master, fcntl.F_SETFL, flags | os.O_NONBLOCK)

    client_sock.setblocking(False)

    try:
        while proc.poll() is None:
            # 使用 select 监听
            rlist, _, _ = select.select([master, client_sock], [], [], 1)

            # 从客户端读取数据
            if client_sock in rlist:
                try:
                    data = client_sock.recv(4096)
                    if not data:
                        break
                    # 写入到 shell
                    os.write(master, data)
                except (BlockingIOError, ConnectionResetError):
                    break

            # 从 shell 读取输出
            if master in rlist:
                try:
                    data = os.read(master, 4096)
                    if not data:
                        break
                    client_sock.sendall(data)
                except (BlockingIOError, OSError):
                    break

    except (OSError, KeyboardInterrupt):
        pass
    finally:
        print(f"[-] 连接结束：{client_addr[0]}:{client_addr[1]}")
        client_sock.close()
        os.close(master)
        if proc.poll() is None:
            proc.terminate()  # 通知进程结束
            try:
                proc.wait(timeout=2) # 等待进程结束
            except subprocess.TimeoutExpired:
                proc.kill()  # 强制结束进程

def main():
    parser = argparse.ArgumentParser(description="TCP to Interactive Shell with Terminal Fix")
    parser.add_argument("-b", "--bind", default="0.0.0.0", help="监听地址 (默认: 0.0.0.0)")
    parser.add_argument("-p", "--port", type=int, default=4444, help="监听端口 (默认: 4444)")
    parser.add_argument("--cols", type=int, default=80, help="初始列数 (默认: 80)")
    parser.add_argument("--rows", type=int, default=24, help="初始行数 (默认: 24)")
    args = parser.parse_args()

    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

    try:
        server.bind((args.bind, args.port))
        server.listen(5)
        print(f"[*] 正在监听 {args.bind}:{args.port} ...")
        print(f"[*] 使用: nc {args.bind} {args.port} 或 socat -,raw,echo=0 tcp:{args.bind}:{args.port}")

        while True:
            try:
                client_sock, client_addr = server.accept()
                handle_client(client_sock, client_addr, initial_cols=args.cols, initial_rows=args.rows)
            except KeyboardInterrupt:
                print("\n[!] 服务已停止")
                break

    except Exception as e:
        print(f"[-] 错误: {e}")
        sys.exit(1)
    finally:
        server.close()

if __name__ == "__main__":
    main()